﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;
using System.ComponentModel;

namespace Microsoft.MSAS.ManagedAPI
{
    // The caller implemented callback for MsasCategorizeUrl().  Will get the 
    // variant hashes in the blob from generic to specific.  You will only get 
    // the interesting hashes (not ones that don't exist in the bloom).  If you return
    // S_OK you must set the bloom and all the category data.  You must return S_FALSE
    // for a cache misses, not to be confused with cached no matches.  For a cached no match,
    // set pbNoMatch to TRUE and return S_OK.  Any other return error will cause enumeration
    // to be terminated and that error to be return from MsasCategorizeUrl().
    //
    //HRESULT CALLBACK
    //MsasCategorizeUrlProc(
    //    __in_bcount(cbHash) const BYTE * pbaHash,  // SHA256 hash of URL variant
    //    __in DWORD cbHash, // always SHA256_DIGEST_LEN (32 bytes)
    //    __out_opt BOOL * pbNoMatch,  // when returning S_OK, set FALSE if setting bloom & cat data or TRUE if not
    //    __out_bcount_opt(cbChildFilter) BYTE * pbaChildFilter,  // bloom data
    //    __in DWORD cbChildFilter,  // always MSAS_CHILDFILTER_SIZE (33 bytes)
    //    __in_bcount(cbCategories) const DWORD * pdwaCategoryIds,  // list of category ids we need data for
    //    __out_bcount(cbCategories) DWORD * pbaCategoryData,  // corresponding list of category data to return (see MSAS_CAT_* values above for more information)
    //    __in DWORD cbCategories,  // number of category ids we need data for
    //    __in_opt LPVOID pvContext  // caller's context data supplied to MsasProcessRequest*()
    //    );
    //typedef HRESULT ( CALLBACK * PFNMSASAPI_CATEGORIZEURLPROC )(
    //    __in_bcount(cbHash) const BYTE * pbaHash,  // SHA256 hash of URL variant
    //    __in DWORD cbHash, // always SHA256_DIGEST_LEN (32 bytes)
    //    __out_opt BOOL * pbNoMatch,  // when returning S_OK, set FALSE if setting bloom & cat data or TRUE if not
    //    __out_bcount_opt(cbChildFilter) BYTE * pbaChildFilter,  // bloom data
    //    __in DWORD cbChildFilter,  // always MSAS_CHILDFILTER_SIZE (33 bytes)
    //    __in_bcount(cbCategories) const DWORD * pdwaCategoryIds,  // list of category ids we need data for
    //    __out_bcount(cbCategories) DWORD * pbaCategoryData,  // corresponding list of category data to return (see MSAS_CAT_* values above for more information)
    //    __in DWORD cbCategories,  // number of category ids we need data for
    //    __in_opt LPVOID pvContext  // caller's context data supplied to MsasProcessRequest*()
    //    );

    [UnmanagedFunctionPointer(CallingConvention.Winapi, SetLastError = false, CharSet = CharSet.Unicode)]
    public delegate uint MrsCategoriezeUrlProc(
        IntPtr pbaHash,  // SHA256 hash of URL variant
        int cbHash, // always SHA256_DIGEST_LEN (32 bytes)
        IntPtr pbNoMatch,  // when returning S_OK, set FALSE if setting bloom & cat data or TRUE if not
        IntPtr pbaChildFilter,  // bloom data
        int cbChildFilter,  // always MSAS_CHILDFILTER_SIZE (33 bytes)
        IntPtr pdwaCategoryIds,  // list of category ids we need data for
        IntPtr pdwaCategoryData,  // corresponding list of category data to return (see MSAS_CAT_* values above for more information)
        int cbCategories,  // number of category ids we need data for
        IntPtr pvContext  // caller's context data supplied to MsasProcessRequest*()
    );

    public class ClientCategorizeProcessor // : IClientApiCategorize
    {
        private const int DEFAULT_DATABLOCK_BUFFER_SIZE = 640;

        #region MRS API
        // Main function that should be used to create the Canonical Url class object.
        //
        //#define MSASAPI_CREATECANONICALURLW "MsasCreateCanonicalUrlW"
        //HRESULT WINAPI
        //MsasCreateCanonicalUrlW(
        //    __in LPCWSTR pszUrl,
        //    __in DWORD dwCanonicalFlags,  // MSAS_CANONICAL_FLAGS
        //    __deref_out MsasApi::IRefCanonicalUrl ** ppCanonicalUrl
        //    );
        //typedef HRESULT ( WINAPI * PFNMSASAPI_CREATECANONICALURLW )(
        //    LPCWSTR,
        //    DWORD,
        //    MsasApi::IRefCanonicalUrl **
        //    );
        [DllImport("msasapi.dll", EntryPoint="MsasCreateCanonicalUrlW", SetLastError = false, CharSet = CharSet.Unicode)]
        public static extern uint MrsCreateCanonicalUrlW(
            string pszUrl,
            MRS_Canonical_Flags dwCanonicalFlags,  // MSAS_CANONICAL_FLAGS
            out UIntPtr ppCanonicalUrl);

        //#define MSASAPI_DESTROYCANONICALURL "MsasDestroyCanonicalUrl"
        //void WINAPI
        //MsasDestroyCanonicalUrl(
        //    __in MsasApi::IRefCanonicalUrl * pUrl
        //    );
        //typedef void ( WINAPI * PFNMSASAPI_DESTROYCANONICALURL )(
        //    __in MsasApi::IRefCanonicalUrl * pUrl
        //    );
        [DllImport("msasapi.dll", EntryPoint = "MsasDestroyCanonicalUrl", SetLastError = false, CharSet = CharSet.Unicode)]
        public static extern uint MrsDestroyCanonicalUrl(
            UIntPtr ppCanonicalUrl);

        // This will use the supplied callback to get data for the URL variants and classify
        // the requested URL.  If all the requried data is present, the function will return S_OK
        // and the pdwCategoryStats will contain the status for each category as defined by the enum
        // MsasApi::CATEGORY_STATUS in msasapi.h.  If there is missing data, S_FALSE is returned and
        // the contents of pdwCategoryStats are undefined.  This is the same for any other FAILED() value
        // as well.
        //
        //#define MSASAPI_CATEGORIZEURL "MsasCategorizeUrl"
        //HRESULT WINAPI
        //MsasCategorizeUrl(
        //    __in MsasApi::IRefCanonicalUrl * pUrl,  // canonical URL object as returned from MsasCreateCanonicalUrlW()
        //    __in_ecount(ceCategories) const DWORD * pdwCategoryIds,  // list of category ids you want classified
        //    __out_ecount(ceCategories) DWORD * pdwCategoryStats,  // corresponding list of category status values (MsasApi::CATEGORY_STATUS from msasapi.h)
        //    __in DWORD ceCategories,  // number of category ids and size of output for category status list
        //    __in DWORD dwFlags,  // MsasApi::CATEGORIZE_FLAGS from msasapi.h
        //    __in PFNMSASAPI_CATEGORIZEURLPROC pFunc,  // caller supplied MsasProcessRequestProc()
        //    __in_opt LPVOID pvContext  // optional caller context data to pass to MsasProcessRequestProc()
        //    );
        //typedef HRESULT ( WINAPI * PFNMSASAPI_CATEGORIZEURL )(
        //    __in MsasApi::IRefCanonicalUrl * pUrl,  // canonical URL object as returned from MsasCreateCanonicalUrlW()
        //    __in_ecount(ceCategories) const DWORD * pdwCategoryIds,  // list of category ids you want classified
        //    __out_ecount(ceCategories) DWORD * pdwCategoryStats,  // corresponding list of category status values (MsasApi::CATEGORY_STATUS from msasapi.h)
        //    __in DWORD ceCategories,  // number of category ids and size of output for category status list
        //    __in DWORD dwFlags,  // MsasApi::CATEGORIZE_FLAGS from msasapi.h
        //    __in PFNMSASAPI_CATEGORIZEURLPROC pFunc,  // caller supplied MsasProcessRequestProc()
        //    __in_opt LPVOID pvContext  // optional caller context data to pass to MsasProcessRequestProc()
        //    );
        [DllImport("msasapi.dll", EntryPoint="MsasCategorizeUrl", SetLastError = false, CharSet = CharSet.Unicode)]
        public static extern uint MrsCategorizeUrl(
            UIntPtr ppCanonicalUrl,
            uint[] pdwCategoryIds,
            int[] pdwCategoryStats,
            int ceCategories,
            MRS_Categorize_Flags dwFlags,
            MrsCategoriezeUrlProc pFunc,
            IntPtr pvContext);

        //#define MSASAPI_GETREQUESTDATA "MsasGetRequestData"
        //HRESULT WINAPI
        //MsasGetRequestData(
        //    __in MsasApi::IRefCanonicalUrl * pUrl,  // canonical URL object as returned from MsasCreateCanonicalUrlW()
        //    __in DWORD dwFlags,
        //    __out_bcount_part_opt(cbData, *pcbDataReturned) BYTE * pbaData,
        //    __in SIZE_T cbData,
        //    __out SIZE_T * pcbDataReturned  // on HRESULT_FROM_WIN32( ERROR_INSUFFICIENT_BUFFER ) contains the length (cbData) required, otherwise contains the number of bytes copied
        //    );
        //typedef HRESULT ( WINAPI * PFNMSASAPI_GETREQUESTDATA )(
        //    __in MsasApi::IRefCanonicalUrl * pUrl,  // canonical URL object as returned from MsasCreateCanonicalUrlW()
        //    __in DWORD dwFlags,
        //    __out_bcount_part(cbData, *pcbDataReturned) BYTE * pbaData,
        //    __in SIZE_T cbData,
        //    __out SIZE_T * pcbDataReturned  // on HRESULT_FROM_WIN32( ERROR_INSUFFICIENT_BUFFER ) contains the length (cbData) required, otherwise contains the number of bytes copied
        //    );
        [DllImport("msasapi.dll", EntryPoint = "MsasGetRequestData", SetLastError = false, CharSet = CharSet.Unicode)]
        public static extern uint MrsGetRequestData(
            UIntPtr ppCanonicalUrl,
            uint dwFlags,
            byte[] pbaData,
            UIntPtr cbData,
            out UIntPtr pcbDataReturned);

        #endregion // MRS API

        [Flags]
        private enum Inheritance_Flags
        {
            Host = 0x10000,
            Path = 0x20000
        };

        [Flags]
        private enum Category_Flags
        {
            Inherit = 0,
            InCategory = 0x1,
            NotInCategory = 0x2,
            Unknown = 0x3,
        };

        public static readonly byte Conf100 = Convert.ToByte(100);
        public static readonly byte Conf50 = Convert.ToByte(50);
        public static readonly byte Conf0 = Convert.ToByte(0);
        public static readonly byte ConfError = Convert.ToByte(255);


        // The pointer to the RefCanonicalUrl that will be used when doing the categorization calls
        private UIntPtr m_ppCanonicalUrl = UIntPtr.Zero;

        /// <summary>
        /// Finalizer that calls dispose
        /// </summary>
        ~ClientCategorizeProcessor()
        {
            Dispose(false);
        }

        public ClientCategorizeProcessor(string url)
        {
            uint hResult = MrsCreateCanonicalUrlW(url,
                MRS_Canonical_Flags.MSAS_Canonical_None, out m_ppCanonicalUrl);

            if (hResult != MRS_HRESULTS.S_OK)
            {
                throw new Win32Exception((int)hResult);
            }
        }

        public byte[] GetDataBlock()
        {
            byte[] buffer;
            IntPtr dataBlockPtr = IntPtr.Zero;
            UIntPtr requiredSizePtr = UIntPtr.Zero;

            // Determine the size of the buffer needed
            uint hResult = MrsGetRequestData(
                m_ppCanonicalUrl,
                0,
                null,
                UIntPtr.Zero,
                out requiredSizePtr);

            if (hResult != MRS_HRESULTS.S_OK)
            {
                throw new Win32Exception((int)hResult);
            }

            // Now get the datablock back
            UInt64 neededBufferSize = (UInt64)requiredSizePtr;
            buffer = new byte[neededBufferSize];
            
            hResult = MrsGetRequestData(
                m_ppCanonicalUrl,
                0,
                buffer,
                (UIntPtr)buffer.Length,
                out requiredSizePtr);

            if (hResult != MRS_HRESULTS.S_OK)
            {
                throw new Win32Exception((int)hResult);
            }

            return buffer;
        }

        public ICollection<SimpleCategory> CategorizeResponseData(SimpleReputationData[] dataToCategorize)
        {
            GCHandle pvContext = new GCHandle();

            List<SimpleCategory> categoryData = new List<SimpleCategory>();

            if (dataToCategorize == null || dataToCategorize.Length == 0)
            {
                return categoryData;
            }

            try
            {
                uint[] categories = GetCategorySubset(dataToCategorize);

                if (categories.Length == 0)
                {
                    // if the top level domain is a no match from the 
                    // database we will have no categories to categorize
                    // on, so no data.  In this case we should just bail
                    // out and return an empty list of categories back to
                    // the caller of the service.
                    return categoryData;
                }

                int[] pdwCategoryStats = new int[categories.Length];

                pvContext = GCHandle.Alloc(dataToCategorize);

                uint hResult = MrsCategorizeUrl(
                    m_ppCanonicalUrl,
                    categories,
                    pdwCategoryStats,
                    categories.Length,
                    MRS_Categorize_Flags.Categorize_Flag_None,
                    CategoriezeUrlCallback,
                    GCHandle.ToIntPtr(pvContext));

                if (hResult != MRS_HRESULTS.S_OK)
                {
                    throw new Win32Exception((int)hResult);
                }

                for (int index = 0; index < pdwCategoryStats.Length; index++)
                {
                    if (pdwCategoryStats[index] != (int)Category_Flags.Unknown)
                    {
                        byte confidence = ConfError;
                        if (pdwCategoryStats[index] == (int)Category_Flags.InCategory)
                        {
                            confidence = Conf100;
                        }
                        else if (pdwCategoryStats[index] == (int)Category_Flags.NotInCategory)
                        {
                            confidence = Conf0;
                        }
                        categoryData.Add(
                            new SimpleCategory() {
                                categoryId = categories[index],
                                confidence = confidence 
                            });
                    }
                }
            }
            finally
            {
                if (pvContext.IsAllocated)
                {
                    pvContext.Free();
                }
            }
            return categoryData;
        }

        private uint[] GetCategorySubset(SimpleReputationData[] dataToCategorize)
        {
            HashSet<uint> hashSet = new HashSet<uint>();

            foreach (SimpleReputationData repuData in dataToCategorize)
            {
                if (repuData.categories == null)
                {
                    // The categories can be null if the SimpleReputationData is 
                    // for a placeholder(dummy) object
                    continue;
                }

                foreach (InputCategory categoryData in repuData.categories)
                {
                    // We are ignoring the boolean return as we do not 
                    // care if it is in the set or not.
                    hashSet.Add(categoryData.categoryId);
                }
            }
            uint[] array = hashSet.ToArray();
            Array.Sort<uint>(array);
            return array;
        }

        internal static uint CategoriezeUrlCallback(
            IntPtr pbaHash,
            int cbHash,
            IntPtr pbNoMatch,
            IntPtr pbaChildFilter,
            int cbChildFilter,
            IntPtr pdwaCategoryIds,
            IntPtr pdwaCategoryData,
            int cbCategories,
            IntPtr pvContext)
        {
            SimpleReputationData[] dataToCategorize = (SimpleReputationData[])GCHandle.FromIntPtr(pvContext).Target;

            // Copy the SHA256 into managed space
            byte[] sha256 = new byte[cbHash];
            Marshal.Copy(pbaHash, sha256, 0, cbHash);

            // Copy the category ids into managed space
            int[] categoryIds = new int[cbCategories];
            Marshal.Copy(pdwaCategoryIds, categoryIds, 0, cbCategories);

            SimpleReputationData simpleRepuData = null;
            for (int dataIndex = 0; dataIndex < dataToCategorize.Length; dataIndex++)
            {
                if (Equals(sha256, dataToCategorize[dataIndex].sha256))
                {
                    simpleRepuData = dataToCategorize[dataIndex];
                    break;
                }
            }

            // We didn't find the raw category data.  This is an 
            // error because we should always find it.
            if (simpleRepuData == null)
            {
                return MRS_HRESULTS.S_FALSE;
            }

            // Set noMatch to true if we have no bloom data and no category data
            if (simpleRepuData.bloom == null 
                && (simpleRepuData.categories == null || simpleRepuData.categories.Length == 0))
            {
                Marshal.WriteInt32(pbNoMatch, 1);
                return MRS_HRESULTS.S_OK;
            }

            // Set the category data on the outbound array
            int[] categoryData = new int[cbCategories];

            // simpleRepuData.categories will be null if there are no categories for
            // the Reputation Object that was used to build the SimpleReputationData
            // This is the case when we have the following XML representation in the 
            // Reputation Object <a:Reputation><a:Data/></a:Reputation>
            if (simpleRepuData.categories != null)
            {
                for (int categoryIndex = 0; categoryIndex < cbCategories; categoryIndex++)
                {
                    int categoryIdToFind = categoryIds[categoryIndex];
                    foreach (InputCategory searchCategory in simpleRepuData.categories)
                    {
                        if (categoryIdToFind == (int)searchCategory.categoryId)
                        {
                            int packedInfo = ClientCategorizeProcessor.PackCategoryData(searchCategory);
                            categoryData[categoryIndex] = packedInfo;
                            break;
                        }
                    }
                }
            }
            Marshal.WriteInt32(pbNoMatch, 0);
            Marshal.Copy(simpleRepuData.bloom, 0, pbaChildFilter, cbChildFilter);
            Marshal.Copy(categoryData, 0, pdwaCategoryData, cbCategories);

            return MRS_HRESULTS.S_OK;
        }

        /// <summary>
        /// Packs the path inheritance, host inheritance, and confidence into an int
        /// so it can be used in an array
        /// </summary>
        /// <param name="searchCategory">The InputCategory object class to pack</param>
        /// <returns>An int with the data packed down</returns>
        internal static int PackCategoryData(InputCategory searchCategory)
        {
            int flaggedDWORD = 0;

            if (searchCategory.confidence == Conf0)
            {
                flaggedDWORD = (int)Category_Flags.NotInCategory;
            }
            else if (searchCategory.confidence == Conf100)
            {
                flaggedDWORD = (int)Category_Flags.InCategory;
            }
            else if (searchCategory.confidence == Conf50)
            {
                flaggedDWORD = (int)Category_Flags.Unknown;
            }

            if (searchCategory.pathInheritance)
            {
                flaggedDWORD |= (int)Inheritance_Flags.Path;
            }
            if (searchCategory.hostInheritance)
            {
                flaggedDWORD |= (int)Inheritance_Flags.Host;
            }

            return flaggedDWORD;
        }

        internal static bool Equals(byte[] expected, byte[] actual)
        {
            if (expected == null && actual == null)
            {
                return true;
            }
            if ((expected == null && actual != null)
                || (expected != null && actual == null)
                || (expected.Length != actual.Length))
            {
                return false;
            }
            return Convert.ToBase64String(expected).Equals(Convert.ToBase64String(actual));
        }


        #region IDisposable Members

        /// <summary>
        /// Public Dispose for proactive disposing of the object. Should be called explicitly,
        /// or indirectly through the using statement.
        /// </summary>
        public void Dispose()
        {
            Dispose(true);

            GC.SuppressFinalize(this);
        }

        #endregion

        /// <summary>
        /// Dispose implementation. 
        /// </summary>
        /// <param name="disposing">If disposing is true, then dispose is being done proactively and
        /// all owned managed and native resources should be disposed. If false, then dispose is being done
        /// as part of garbage collection, so only native resourses should be disposed.</param>
        protected void Dispose(bool disposing)
        {
            if (m_ppCanonicalUrl != UIntPtr.Zero)
            {
                MrsDestroyCanonicalUrl(m_ppCanonicalUrl);
                m_ppCanonicalUrl = UIntPtr.Zero;
            }
        }
    }

    public class SimpleReputationData
    {
        public byte[] sha256;
        public byte[] bloom;
        public InputCategory[] categories;
    }

    public struct InputCategory
    {
        public uint categoryId;
        public bool pathInheritance;
        public bool hostInheritance;
        public byte confidence;        
    }

    public struct SimpleCategory
    {
        public uint categoryId;
        public byte confidence;
    }
}